//! 负载均衡模块
//! 
//! 负责集群中请求的分发和负载均衡，支持多种负载均衡策略。

use crate::{Error, Result};
use super::{ClusterConfig, ClusterEvent};
use std::collections::HashMap;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::time::SystemTime;
use serde::{Serialize, Deserialize};
use tokio::sync::mpsc;

/// 负载均衡器
#[derive(Debug)]
pub struct LoadBalancer {
    /// 集群配置
    config: ClusterConfig,
    /// 负载均衡策略
    strategy: LoadBalancingStrategy,
    /// 节点权重
    node_weights: HashMap<String, f64>,
    /// 节点负载统计
    node_loads: HashMap<String, NodeLoad>,
    /// 轮询计数器
    round_robin_counter: AtomicUsize,
    /// 事件发送器
    event_sender: Option<mpsc::UnboundedSender<ClusterEvent>>,
}

/// 负载均衡策略
#[derive(Debug, Clone)]
pub enum LoadBalancingStrategy {
    /// 轮询
    RoundRobin,
    /// 加权轮询
    WeightedRoundRobin,
    /// 最少连接
    LeastConnections,
    /// 加权最少连接
    WeightedLeastConnections,
    /// 响应时间
    ResponseTime,
    /// 一致性哈希
    ConsistentHash,
}

/// 节点负载信息
#[derive(Debug, Clone)]
pub struct NodeLoad {
    /// 节点ID
    pub node_id: String,
    /// 活跃连接数
    pub active_connections: usize,
    /// 平均响应时间（毫秒）
    pub avg_response_time_ms: f64,
    /// CPU使用率
    pub cpu_usage: f64,
    /// 内存使用率
    pub memory_usage: f64,
    /// 请求计数
    pub request_count: u64,
    /// 最后更新时间
    pub last_updated: SystemTime,
}

/// 负载均衡结果
#[derive(Debug, Clone)]
pub struct BalancingResult {
    /// 选中的节点ID
    pub selected_node: String,
    /// 选择原因
    pub reason: String,
    /// 负载分数
    pub load_score: f64,
}

impl LoadBalancer {
    /// 创建新的负载均衡器
    pub fn new(config: ClusterConfig) -> Result<Self> {
        Ok(Self {
            config,
            strategy: LoadBalancingStrategy::RoundRobin,
            node_weights: HashMap::new(),
            node_loads: HashMap::new(),
            round_robin_counter: AtomicUsize::new(0),
            event_sender: None,
        })
    }
    
    /// 启动负载均衡器
    pub async fn start(&mut self) -> Result<()> {
        tracing::info!("启动负载均衡器");
        
        // 初始化当前节点的负载信息
        self.initialize_node_loads().await?;
        
        tracing::info!("负载均衡器启动完成");
        Ok(())
    }
    
    /// 停止负载均衡器
    pub async fn stop(&mut self) -> Result<()> {
        tracing::info!("停止负载均衡器");
        Ok(())
    }
    
    /// 初始化节点负载信息
    async fn initialize_node_loads(&mut self) -> Result<()> {
        let node_load = NodeLoad {
            node_id: self.config.node_id.clone(),
            active_connections: 0,
            avg_response_time_ms: 0.0,
            cpu_usage: 0.0,
            memory_usage: 0.0,
            request_count: 0,
            last_updated: SystemTime::now(),
        };
        
        self.node_loads.insert(self.config.node_id.clone(), node_load);
        self.node_weights.insert(self.config.node_id.clone(), 1.0);
        
        Ok(())
    }
    
    /// 选择节点
    pub async fn select_node(&mut self, request_key: Option<&str>) -> Result<BalancingResult> {
        let available_nodes: Vec<String> = self.node_loads.keys().cloned().collect();
        
        if available_nodes.is_empty() {
            return Err(Error::load_balancing("没有可用的节点".to_string()));
        }
        
        let result = match self.strategy {
            LoadBalancingStrategy::RoundRobin => {
                self.round_robin_select(&available_nodes)
            }
            LoadBalancingStrategy::WeightedRoundRobin => {
                self.weighted_round_robin_select(&available_nodes)
            }
            LoadBalancingStrategy::LeastConnections => {
                self.least_connections_select(&available_nodes)
            }
            LoadBalancingStrategy::WeightedLeastConnections => {
                self.weighted_least_connections_select(&available_nodes)
            }
            LoadBalancingStrategy::ResponseTime => {
                self.response_time_select(&available_nodes)
            }
            LoadBalancingStrategy::ConsistentHash => {
                self.consistent_hash_select(&available_nodes, request_key)
            }
        };
        
        tracing::debug!("选择节点: {:?}", result);
        Ok(result)
    }
    
    /// 轮询选择
    fn round_robin_select(&self, nodes: &[String]) -> BalancingResult {
        let index = self.round_robin_counter.fetch_add(1, Ordering::Relaxed) % nodes.len();
        let selected_node = nodes[index].clone();
        
        BalancingResult {
            selected_node,
            reason: "轮询选择".to_string(),
            load_score: 1.0,
        }
    }
    
    /// 加权轮询选择
    fn weighted_round_robin_select(&self, nodes: &[String]) -> BalancingResult {
        // 简化实现：使用权重最高的节点
        let mut best_node = nodes[0].clone();
        let mut best_weight = 0.0;
        
        for node in nodes {
            if let Some(&weight) = self.node_weights.get(node) {
                if weight > best_weight {
                    best_weight = weight;
                    best_node = node.clone();
                }
            }
        }
        
        BalancingResult {
            selected_node: best_node,
            reason: format!("加权轮询选择，权重: {}", best_weight),
            load_score: best_weight,
        }
    }
    
    /// 最少连接选择
    fn least_connections_select(&self, nodes: &[String]) -> BalancingResult {
        let mut best_node = nodes[0].clone();
        let mut min_connections = usize::MAX;
        
        for node in nodes {
            if let Some(load) = self.node_loads.get(node) {
                if load.active_connections < min_connections {
                    min_connections = load.active_connections;
                    best_node = node.clone();
                }
            }
        }
        
        BalancingResult {
            selected_node: best_node,
            reason: format!("最少连接选择，连接数: {}", min_connections),
            load_score: min_connections as f64,
        }
    }
    
    /// 加权最少连接选择
    fn weighted_least_connections_select(&self, nodes: &[String]) -> BalancingResult {
        let mut best_node = nodes[0].clone();
        let mut best_score = f64::MAX;
        
        for node in nodes {
            if let (Some(load), Some(&weight)) = (self.node_loads.get(node), self.node_weights.get(node)) {
                let score = if weight > 0.0 {
                    load.active_connections as f64 / weight
                } else {
                    f64::MAX
                };
                
                if score < best_score {
                    best_score = score;
                    best_node = node.clone();
                }
            }
        }
        
        BalancingResult {
            selected_node: best_node,
            reason: format!("加权最少连接选择，分数: {:.2}", best_score),
            load_score: best_score,
        }
    }
    
    /// 响应时间选择
    fn response_time_select(&self, nodes: &[String]) -> BalancingResult {
        let mut best_node = nodes[0].clone();
        let mut min_response_time = f64::MAX;
        
        for node in nodes {
            if let Some(load) = self.node_loads.get(node) {
                if load.avg_response_time_ms < min_response_time {
                    min_response_time = load.avg_response_time_ms;
                    best_node = node.clone();
                }
            }
        }
        
        BalancingResult {
            selected_node: best_node,
            reason: format!("响应时间选择，响应时间: {:.2}ms", min_response_time),
            load_score: min_response_time,
        }
    }
    
    /// 一致性哈希选择
    fn consistent_hash_select(&self, nodes: &[String], request_key: Option<&str>) -> BalancingResult {
        let key = request_key.unwrap_or("default");
        let hash = self.hash_key(key);
        let index = hash % nodes.len();
        let selected_node = nodes[index].clone();
        
        BalancingResult {
            selected_node,
            reason: format!("一致性哈希选择，键: {}", key),
            load_score: hash as f64,
        }
    }
    
    /// 哈希函数
    fn hash_key(&self, key: &str) -> usize {
        // 简化的哈希函数
        key.bytes().map(|b| b as usize).sum()
    }
    
    /// 更新节点负载
    pub async fn update_node_load(&mut self, node_id: &str, load: NodeLoad) -> Result<()> {
        tracing::debug!("更新节点负载: {} -> {:?}", node_id, load);
        self.node_loads.insert(node_id.to_string(), load);
        Ok(())
    }
    
    /// 设置节点权重
    pub async fn set_node_weight(&mut self, node_id: &str, weight: f64) -> Result<()> {
        tracing::debug!("设置节点权重: {} -> {}", node_id, weight);
        self.node_weights.insert(node_id.to_string(), weight);
        Ok(())
    }
    
    /// 添加节点
    pub async fn add_node(&mut self, node_id: &str, weight: f64) -> Result<()> {
        tracing::info!("添加节点到负载均衡器: {}", node_id);
        
        let node_load = NodeLoad {
            node_id: node_id.to_string(),
            active_connections: 0,
            avg_response_time_ms: 0.0,
            cpu_usage: 0.0,
            memory_usage: 0.0,
            request_count: 0,
            last_updated: SystemTime::now(),
        };
        
        self.node_loads.insert(node_id.to_string(), node_load);
        self.node_weights.insert(node_id.to_string(), weight);
        
        Ok(())
    }
    
    /// 移除节点
    pub async fn remove_node(&mut self, node_id: &str) -> Result<()> {
        tracing::info!("从负载均衡器移除节点: {}", node_id);
        
        self.node_loads.remove(node_id);
        self.node_weights.remove(node_id);
        
        Ok(())
    }
    
    /// 重新平衡
    pub async fn rebalance(&mut self) -> Result<()> {
        tracing::info!("执行负载重新平衡");
        
        // 重新计算节点权重
        self.recalculate_weights().await?;
        
        // 发送重新平衡事件
        if let Some(sender) = &self.event_sender {
            let _ = sender.send(ClusterEvent::ShardRebalanced);
        }
        
        Ok(())
    }
    
    /// 重新计算权重
    async fn recalculate_weights(&mut self) -> Result<()> {
        for (node_id, load) in &self.node_loads {
            // 基于CPU和内存使用率计算权重
            let cpu_factor = 1.0 - (load.cpu_usage / 100.0);
            let memory_factor = 1.0 - (load.memory_usage / 100.0);
            let weight = (cpu_factor + memory_factor) / 2.0;
            
            self.node_weights.insert(node_id.clone(), weight.max(0.1));
        }
        
        Ok(())
    }
    
    /// 设置负载均衡策略
    pub fn set_strategy(&mut self, strategy: LoadBalancingStrategy) {
        tracing::info!("设置负载均衡策略: {:?}", strategy);
        self.strategy = strategy;
    }
    
    /// 获取节点负载信息
    pub fn get_node_load(&self, node_id: &str) -> Option<&NodeLoad> {
        self.node_loads.get(node_id)
    }
    
    /// 获取所有节点负载
    pub fn get_all_node_loads(&self) -> Vec<&NodeLoad> {
        self.node_loads.values().collect()
    }
    
    /// 设置事件发送器
    pub fn set_event_sender(&mut self, sender: mpsc::UnboundedSender<ClusterEvent>) {
        self.event_sender = Some(sender);
    }
}

impl crate::Error {
    /// 创建负载均衡错误
    pub fn load_balancing<S: Into<String>>(msg: S) -> Self {
        Self::Storage(msg.into())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    
    #[tokio::test]
    async fn test_load_balancer_creation() {
        let config = ClusterConfig::default();
        let load_balancer = LoadBalancer::new(config).unwrap();
        
        assert!(matches!(load_balancer.strategy, LoadBalancingStrategy::RoundRobin));
        assert_eq!(load_balancer.node_weights.len(), 0);
        assert_eq!(load_balancer.node_loads.len(), 0);
    }
    
    #[tokio::test]
    async fn test_node_management() {
        let config = ClusterConfig::default();
        let mut load_balancer = LoadBalancer::new(config).unwrap();
        
        // 添加节点
        load_balancer.add_node("node-1", 1.0).await.unwrap();
        load_balancer.add_node("node-2", 2.0).await.unwrap();
        
        assert_eq!(load_balancer.node_loads.len(), 2);
        assert_eq!(load_balancer.node_weights.len(), 2);
        
        // 移除节点
        load_balancer.remove_node("node-1").await.unwrap();
        
        assert_eq!(load_balancer.node_loads.len(), 1);
        assert!(!load_balancer.node_loads.contains_key("node-1"));
    }
    
    #[tokio::test]
    async fn test_round_robin_selection() {
        let config = ClusterConfig::default();
        let mut load_balancer = LoadBalancer::new(config).unwrap();
        
        // 添加节点
        load_balancer.add_node("node-1", 1.0).await.unwrap();
        load_balancer.add_node("node-2", 1.0).await.unwrap();
        
        // 测试轮询选择
        let result1 = load_balancer.select_node(None).await.unwrap();
        let result2 = load_balancer.select_node(None).await.unwrap();
        
        // 应该选择不同的节点（轮询）
        assert_ne!(result1.selected_node, result2.selected_node);
    }
    
    #[tokio::test]
    async fn test_consistent_hash_selection() {
        let config = ClusterConfig::default();
        let mut load_balancer = LoadBalancer::new(config).unwrap();
        load_balancer.set_strategy(LoadBalancingStrategy::ConsistentHash);
        
        // 添加节点
        load_balancer.add_node("node-1", 1.0).await.unwrap();
        load_balancer.add_node("node-2", 1.0).await.unwrap();
        
        // 测试一致性哈希选择
        let result1 = load_balancer.select_node(Some("key1")).await.unwrap();
        let result2 = load_balancer.select_node(Some("key1")).await.unwrap();
        
        // 相同的键应该选择相同的节点
        assert_eq!(result1.selected_node, result2.selected_node);
    }
    
    #[tokio::test]
    async fn test_load_balancer_rebalance() {
        let config = ClusterConfig::default();
        let mut load_balancer = LoadBalancer::new(config).unwrap();
        
        // 添加节点
        load_balancer.add_node("node-1", 1.0).await.unwrap();
        
        // 更新节点负载
        let node_load = NodeLoad {
            node_id: "node-1".to_string(),
            active_connections: 10,
            avg_response_time_ms: 100.0,
            cpu_usage: 50.0,
            memory_usage: 60.0,
            request_count: 1000,
            last_updated: SystemTime::now(),
        };
        
        load_balancer.update_node_load("node-1", node_load).await.unwrap();
        
        // 执行重新平衡
        load_balancer.rebalance().await.unwrap();
        
        // 验证权重已更新
        let weight = load_balancer.node_weights.get("node-1").unwrap();
        assert!(*weight > 0.0 && *weight <= 1.0);
    }
}
